# -*- coding: utf-8 -*-

# 导入必要的工具包
import xgboost as xgb

# 计算分类正确率
from sklearn.metrics import accuracy_score

# read in data，数据在xgboost安装的路径下的demo目录,现在我们将其copy到当前代码下的data目录
dpath = '../course/data/'
dtrain = xgb.DMatrix(dpath + 'agaricus.txt.train')
dtest = xgb.DMatrix(dpath + 'agaricus.txt.test')

print(dtrain.num_col())
print(dtrain.num_row())

# specify parameters via map
param = {'max_depth':2, 'eta':1, 'silent':0, 'objective':'binary:logistic' }

# 设置boosting迭代计算次数，即若学习器（决策树）的数目
num_round = 2

import time
starttime = time.clock()

bst = xgb.train(param, dtrain, num_round)

endtime = time.clock()
print (endtime - starttime)

train_preds = bst.predict(dtrain)
train_predictions = [round(value) for value in train_preds]
y_train = dtrain.get_label()
train_accuracy = accuracy_score(y_train, train_predictions)
print ("Train Accuary: %.2f%%" % (train_accuracy * 100.0))

# make prediction
test_preds = bst.predict(dtest)
test_predictions = [round(value) for value in test_preds]
y_test = dtest.get_label()
test_accuracy = accuracy_score(y_test, test_predictions)
print ("Train Accuary: %.2f%%" % (test_accuracy * 100.0))

from matplotlib import pyplot
import graphviz
#xgb.plot_tree(bst, num_trees=0, rankdir= 'LR' )
#pyplot.show()

#xgb.plot_tree(bst,num_trees=1, rankdir= 'LR' )
#pyplot.show()
#xgb.to_graphviz(bst,num_trees=0)
xgb.to_graphviz(bst,num_trees=1)